# 不经过第一步粗分类，直接使用combine预料进行第二步的推理
from utils.multiprocessing_inf import start_infrance, getcorp
import pickle

if __name__ == '__main__':
    # --用二分类数据集测试分布式推理有无BUG---
    # testcorp = test4bin_3_overall()
    # corp = []
    # for each in testcorp:
    #     corp.append(each[0].split())
    # --------------------------------------
    corp = pickle.load(open("./corps/combinecorp_idx.pt", "rb"))
    parts = 300  # 分300块处理合成语料,因为有三块卡，所以设置成三的倍数
    blocksize = len(corp) // parts
    corp_part_li = {}
    for part in range(parts):
        corp_part = corp[part * blocksize:(part + 1) * blocksize]
        corp_part = getcorp(corp_part)
        corp_part_li[part % 3] = corp_part
        print(f"corp part {part} prepared.")
        if part % 3 == 2:
            start_infrance(3, corp_part_li, part)
    # -------------收集结果------------------
    collect = []
    for p in range(parts):
        f = pickle.load(open(f"./corps/pt/infrance_res_part{p}.pt", "rb"))
        collect += f
    # pickle.dump(collect, open(f"./corps/pt/infrance_res.pt", "wb"))  # 800M 1091349个
    pickle.dump(collect, open(f"./corps/pt/infrance_res_!.pt", "wb"))

# nohup python -u bint5_output.py &
